#行时会检查必要的模块是否已经安装，如果没有，则会自动安装所需的模块。
import pkg_resources
REQUIRED_PACKAGES = [
    'colorama'
]

for package in REQUIRED_PACKAGES:
    try:
        dist = pkg_resources.get_distribution(package)
        print('{} ({}) 已安装'.format(dist.key, dist.version))
    except pkg_resources.DistributionNotFound:
        print('{} 尚未安装. 正在自动安装...'.format(package))
        import pip
        pip.main(['install', package])

import os
import sys
import warnings
import subprocess
from colorama import Fore
from concurrent.futures import ThreadPoolExecutor
import re
import rosbag
import itertools
import def_function
import pandas as pd
import os
import json
import wget

# 获取当前脚本的完整路径
script_path = os.path.abspath(sys.argv[0])
# 从完整路径中获取目录
script_dir = os.path.dirname(script_path)
# 从完整路径中分离出文件名
script_name = os.path.basename(script_path)
# 使用 splitext() 函数分离文件名和扩展名
script_name_without_extension, _ = os.path.splitext(script_name)

# print(Fore.YELLOW+f"python script name:", sys.argv[0])
# if len(sys.argv)!=2:
#     print(Fore.RED+f"you need python params!")
#     sys.exit()

# python_param = sys.argv[1]
python_param = '/mnt/08_bag/rocky/outputPL077159_recording_dfdi_console_manual_recording_20230814-120005_20230814-120110_0_augmented.bag'
python_param = 'f42e23c6aa95fd45c3d82aa50f189e5a'
print(Fore.WHITE+f"python_param:", python_param)

base_name = os.path.basename(python_param)
def_function.get_data_info.get_input(base_name)
info_json = def_function.get_data_info.get_info_json()
business_type = info_json.get('business_type')
bag_name = info_json.get('bag_name')
bag_md5 = info_json.get('bag_md5')

dataspace = "./dataspace"
dataspace_current = os.path.abspath(dataspace) + "/" + business_type
if not os.path.exists(dataspace_current):
    os.makedirs(dataspace_current)

# 获取路径下的所有.py文件
scripts = [os.path.join(script_dir, f) for f in os.listdir(script_dir) if script_name not in f and f.endswith(".py")]
# 使用 os.walk 来递归获取所有 .py 文件
# scripts = [os.path.join(root, f) 
#            for root, dirs, files in os.walk(current_path) 
#            for f in files if f.endswith(".py")]

# 打开文件并开始读取
topics_list = list()
print(Fore.GREEN+f"===============================================================")
for script in scripts:
    print(Fore.GREEN+f"", os.path.basename(script))
    with open(script, 'r') as file:
        # 遍历每一行
        for line in file:
            if line.strip().startswith('topics = ['):
                topics_list.append(eval(line.split('=')[1].strip()))
            # 如果该行包含 'error_type ='
            if 'error_type =' in line :
                match = re.search(r'error_type = "(.*)"', line)
                # 如果找到匹配的结果
                if match:
                    error_and_condition = match.group(1)
            if 'judge_condition =' in line :
                match = re.search(r'judge_condition = "(.*)"', line)
                if match:
                    error_and_condition = error_and_condition + "{" + match.group(1) + "}"
                    # 打印出双引号中的内容
                    print(Fore.MAGENTA+f"\t", error_and_condition)

# 打印出当前所有所需要的topics
# 使用 itertools.chain.from_iterable 展平嵌套列表
flattened_topics_list = list(itertools.chain.from_iterable(topics_list))
topics_set = set(flattened_topics_list)
print(Fore.YELLOW+f"===============================================================")
for topic in topics_set:
    print(Fore.YELLOW+f"topic", topic)

current_rosbag = dataspace_current + "/" + bag_name;
filter_rosbag = current_rosbag.replace(".bag", "_filter.bag")
temp_rosbag = current_rosbag.replace(".bag", "_temp.bag")
if not os.path.exists(filter_rosbag):
    if not os.path.isabs(python_param):
        #非绝对路径，即期望网络下载包分析
        if not os.path.exists(current_rosbag):
            url = 'https://data.momenta.works/api/v2/data/uri?id='+bag_md5+'&https=1'
            wget.download(url, out = current_rosbag)
    # filter_command = "rosbag filter " + current_rosbag + " " + temp_rosbag
    # topic_param = '"'
    # for index, topic in enumerate(topics_set):
    #     if index == 0:
    #         topic_param += "topic == '" + topic + "'"
    #     else:
    #         topic_param += " or topic == '" + topic + "'"
    # topic_param += '"'
    # filter_command += " " + topic_param
    # os.system(filter_command)

    # python convert is spent to long time
    print(Fore.GREEN+f"convert rosbag[" + Fore.CYAN+f"" + current_rosbag + Fore.RED+f"=>" + Fore.CYAN+f"" + filter_rosbag + Fore.GREEN+f"]")
    bag = rosbag.Bag(current_rosbag)
    total_duration = bag.get_end_time() - bag.get_start_time()
    with rosbag.Bag(temp_rosbag, 'w') as outbag:
        for topic, msg, t in bag.read_messages():
            try:
                if topic in topics_set:
                    outbag.write(topic, msg, t)
            except UnicodeDecodeError:  # 如果发生 UnicodeDecodeError 则执行以下操作
                print(f"Encountered a UnicodeDecodeError at {topic}, but continuing...")
                continue  # 跳过当前循环，进入下一次循环
            progress = max(0, (t.to_sec() - bag.get_start_time()) / total_duration)
            # Calculate the number of dots and spaces for the progress bar
            dots = "." * int(progress * 50)
            spaces = " " * (50 - len(dots))
            # Generate the progress string
            progress_str = f"[{dots}{spaces}] {t.to_sec()} / {bag.get_end_time()}"
            # Print the progress with carriage return at the end, so it overwrites the current line
            sys.stdout.write(f"\r{Fore.GREEN}{progress * 100:.2f}% {progress_str}")
            sys.stdout.flush()
    os.rename(temp_rosbag, filter_rosbag)

# 定义用于执行脚本的函数
def run_script(scripts):
    #print(Fore.WHITE+f"python3", scripts, current_rosbag)
    subprocess.run(["python3", scripts, filter_rosbag])

# 使用线程池来并行运行脚本
with ThreadPoolExecutor() as executor:
    if script_dir != scripts: #不运行自身
        pass
        executor.map(run_script, scripts)